# -*- coding: utf-8 -*-
"""
Created on Sun Oct 22 10:12:04 2017
@author: Yuxi Chi 
@email: cherichy@outlook.com
"""
import sys
import numpy as np
from scipy import linalg
import matplotlib.pyplot as plt

class Material:
    def __init__(self, E=2e11, mu=0.3, rho=7800):
        self.E = E
        self.mu = mu
        self.rho = rho

class Section:
    def __init__(self, A=0.01):
        self.A = A

class Element:
    def __init__(self, material: Material, section: Section, data: np.ndarray):
        self.E = material.E
        self.rho = material.rho
        self.A = section.A
        self.data = data

def chop(x):
    x[abs(x) < abs(x).max() / 10 ** 10] = 0  # 相差10个量级即舍去

class Truss2D:
    def __init__(self):
        self.materials = {}
        self.sections = {}
        self.elements = []

    def paser_inp(self, inpname):
        self.inpname = inpname
        inp = open(inpname)
        inpdata = inp.read()
        inp.close()
        for part in inpdata.split('*'):
            line = part[:-1].split('\n')
            # Material
            if line[0].startswith('Material'):
                material_name = line[0].split('=')[-1]
                self.materials[material_name] = eval('Material(' + ','.join(line[1:]) + ')')
            # Section
            if line[0].startswith('Section'):
                section_name = line[0].split('=')[-1]
                self.sections[section_name] = eval('Section(' + ','.join(line[1:]) + ')')
            # Node
            if line[0].startswith('Node'):
                self.nodes = np.array(np.matrix(';'.join(line[1:])), dtype=np.float64)
                size = len(self.nodes)
                self.KK = np.zeros((size * 2, size * 2), dtype=np.float64)  # Global KK Matrix
                self.MM = np.zeros((size * 2, size * 2), dtype=np.float64)
                self.UU = np.zeros(size * 2, dtype=np.float64)  # Displacement Vector
                self.FF = np.zeros(size * 2, dtype=np.float64)  # Force Vector
                self.bc_unknow = np.ones(size * 2, dtype=bool)
            # Element
            if line[0].startswith('Element'):
                material_name = line[0].split(',')[1].split('=')[1]
                section_name = line[0].split(',')[2].split('=')[1]
                self.elements.append(
                    Element(self.materials[material_name],
                            self.sections[section_name],
                            np.array(np.matrix(';'.join(line[1:])))))
            # BC
            if line[0].startswith('BC'):
                for bcs in line[1:]:
                    node, dof, disp = eval(bcs)
                    # only consider 0 BC (not specific displacement)
                    self.UU[2 * node - 3 + dof] = disp  # 1,0,1=>0  2,0,2=>3
                    self.bc_unknow[2 * node - 3 + dof] = False
            # Force
            if line[0].startswith('Force'):
                for bcs in line[1:]:
                    node, dof, force = eval(bcs)
                    self.FF[2 * node - 3 + dof] = force


    def assemble_global_K_M(self):
        for element in self.elements:
            for rod in element.data:
                n1, n2 = self.nodes[(rod - 1)[1:]]  # rod: index,n1,n2
                L = linalg.norm((n2 - n1)[1:])  # n: index,x,y
                K = element.E * element.A / L * np.array([[1, -1], [-1, 1]])
                # 一致质量矩阵
                # M = element.rho * element.A * L / 6 * np.array([[2, 1], [1, 2]])
                # 堆聚质量矩阵
                M = element.rho * element.A * L / 2 * np.array([[1, 0], [0, 1]])
                # arctan2(dy,dx),*将数组解包作为参数
                alpha = np.arctan2(*((n2 - n1)[-1:0:-1]))  
                C = np.cos(alpha)
                S = np.sin(alpha)
                # 坐标转换矩阵
                Trans = np.array([[C, S, 0, 0], [0, 0, C, S]])
                # @ : matmul
                GK = Trans.T @ K @ Trans  
                GM = Trans.T @ M @ Trans

                i, j = int(n1[0]), int(n2[0])  # i,j : node index
                index = [2 * i - 2, 2 * i - 1, 2 * j - 2, 2 * j - 1]
                self.KK[np.ix_(index, index)] += GK
                self.MM[np.ix_(index, index)] += GM

    def solve_static(self):
        #  [K11,K12] [U1] [F1]    U1: unknown BC
        #  [K12,K22] [U2] [F2]    U2: known BC
        #  K11*U1+K12*U2=F1 ==> U1=K11'*(F1-K12*U2)
        #  [F1,F2]=FF+RF ==> RF = K*U-FF

        # np.ix_ : 取出对应行列的子矩阵
        K11 = self.KK[np.ix_(self.bc_unknow, self.bc_unknow)]
        K12 = self.KK[np.ix_(self.bc_unknow, np.logical_not(self.bc_unknow))]
        U2 = self.UU[np.logical_not(self.bc_unknow)]
        F1 = self.FF[self.bc_unknow]
        # Solve U1
        self.UU[self.bc_unknow] = linalg.solve(K11, F1 - K12 @ U2)
        # Get reaction force
        self.RF = self.KK @ self.UU - self.FF
        # Element Stress
        self.SS = []
        for element in self.elements:
            self.SS.append(np.zeros((len(element.data), 3)))
            for index, ele in enumerate(element.data):
                n1, n2 = self.nodes[(ele - 1)[1:]]
                alpha = np.arctan2(*((n2 - n1)[-1:0:-1]))
                C = np.cos(alpha)
                S = np.sin(alpha)
                i, j = int(n1[0]), int(n2[0])
                strain = np.array([-C, -S, C, S]) @ self.UU[
                    [2 * i - 2, 2 * i - 1, 2 * j - 2, 2 * j - 1]] / linalg.norm(
                    (n2 - n1)[1:])
                stress = element.E * strain
                inn_force = stress * element.A
                self.SS[-1][index] = stress, strain, inn_force

    def output_static(self):
        chop(self.UU)
        for i in self.SS:
            chop(i[:, 0])
            chop(i[:, 1])
            chop(i[:, 2])
        chop(self.RF)

        with open(self.inpname.split('.')[-2] + '_static.txt', 'w') as f:
            f.write("File:\t%*s\n" % (10, filename))
            f.write('\n' + '-' * 50 + '\n')
            f.write(open(self.inpname).read())
            f.write('\n' + '-' * 50 + '\n')
            f.write("\nResult:\n")
            f.write("\n*Node Displacement:\n")
            f.write("  Node\t|%*s\t%*s\t%*s\n" % (8, 'X', 8, 'Y', 10, 'VSum'))
            f.write('-' * 55 + '\n')
            for i in range(len(self.nodes)):
                if self.UU[2 * i] != 0 or self.UU[2 * i + 1] != 0:
                    f.write("%4d\t| % 10E\t% 10E\t% 10E\n" % (
                        i + 1, self.UU[2 * i], self.UU[2 * i + 1],
                        linalg.norm([self.UU[2 * i], self.UU[2 * i + 1]])))

            f.write("\n*Reaction Force:\n")
            f.write("  Node\t|%*s\t%*s\t%*s\n" % (8, 'X', 8, 'Y', 10, 'VSum'))
            f.write('-' * 55 + '\n')
            for i in range(len(self.nodes)):
                if self.RF[2 * i] != 0 or self.RF[2 * i + 1] != 0:
                    f.write("%4d\t| % 10E\t% 10E\t% 10E\n" % (
                        i + 1, self.RF[2 * i], self.RF[2 * i + 1],
                        linalg.norm([self.RF[2 * i], self.RF[2 * i + 1]])))

            f.write("\n*Element:\n")
            f.write("Element\t|%*s\t%*s\t%*s\n" % (10, 'Stress', 10, 'Strain', 12, 'InnForce'))
            f.write('-' * 55 + '\n')
            for index in range(len(self.elements)):
                s = self.SS[index]
                for i in range(len(s)):
                    f.write("%4d\t| % 10E\t% 10E\t% 10E\n" % (i + 1, s[i, 0], s[i, 1], s[i, 2]))

    def solve_modal(self):
        eigval, eigvec = linalg.eig(self.KK[np.ix_(self.bc_unknow, self.bc_unknow)],
                                    self.MM[np.ix_(self.bc_unknow, self.bc_unknow)])
        # eigval从小到大排序
        eigvec = eigvec[:, eigval.argsort()]
        eigval.sort()
        self.freq = np.sqrt(np.real(eigval))
        self.modvec = np.zeros((len(self.nodes) * 2, len(eigval)), dtype=np.float64)
        for i in range(len(eigval)):
            # eigvec归一化，除以每列的最大值
            eigvec[:, i] /= np.abs(eigvec[np.abs(eigvec[:, i]).argmax(), i])
            chop(eigvec[:, i])
            self.modvec[self.bc_unknow, i] = eigvec[:, i]

    def output_modal(self):
        with open(self.inpname.split('.')[-2] + '_modal.txt', 'w') as f:
            f.write("File:\t%*s\n" % (10, filename))
            f.write('\n' + '-' * 50 + '\n')
            f.write(open(self.inpname).read())
            f.write('\n' + '-' * 50 + '\n')
            f.write("\nResult:\n")
            for index in range(len(self.freq)):
                f.write("*Modal %d, Freq =% 10E Hz\n" % (index + 1, self.freq[index]))
                f.write("*Modal Vector:\n")
                for i, v in enumerate(self.modvec[:, index]):
                    f.write("%d\t% 10E\n" % (i + 1, v))
                f.write("\n")

    def show_structure(self, defvec: np.ndarray = None, title: str = 'Deformation', show_ref: bool = True):
        plt.figure()
        def_loc = self.nodes[:, 1:].copy()
        # 归一化以便显示
        # node_range=def_loc.max()
        defvec = 0.3 * defvec / defvec.max()
        if show_ref:
            for element in self.elements:
                for rod in element.data:
                    n1d, n2d = def_loc[(rod - 1)[1:]]
                    plt.plot([n1d[0], n2d[0]], [n1d[1], n2d[1]], 'k.-', alpha=0.6)
        if not defvec is None:
            def_loc += defvec.reshape((-1, 2))
            for element in self.elements:
                for rod in element.data:
                    n1d, n2d = def_loc[(rod - 1)[1:]]
                    plt.plot([n1d[0], n2d[0]], [n1d[1], n2d[1]], 'co-', alpha=0.6)
        plt.title(title)
        plt.axis('equal')
        plt.show()

    def show_modal(self, index):
        defvec = self.modvec[:, index - 1]
        freq = self.freq[index - 1]
        self.show_structure(defvec, 'Modal %d,Freq=%.4fHz' % (index, freq))

def main(filename):
    Job = Truss2D()
    Job.paser_inp(filename)
    Job.assemble_global_K_M()
    Job.solve_static()
    Job.output_static()
    Job.solve_modal()
    Job.output_modal()
    # Job.show_structure(Job.UU)
    for i in range(len(Job.freq)):
        Job.show_modal(i + 1)

if __name__ == '__main__':
    for filename in sys.argv[1:]:
        main(filename)
